import os

import bs4
from dotenv import load_dotenv
from langchain import hub
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains.retrieval import create_retrieval_chain
from langchain_chroma import Chroma
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

from com.wp.langchain.rag.LLMInit import LLMInit

# 加载环境变量
load_dotenv()
llm = LLMInit().get_llm()

"""
我们需要首先加载博客文章的内容。我们可以使用 文档加载器 来实现这一点，它们是从源加载数据并返回 列表的 文档。 一个 Document 是一个包含一些 page_content (str) 和 metadata (dict) 的对象。

在这种情况下，我们将使用 WebBaseLoader， 使用 urllib 从网页 URL 加载 HTML，并使用 BeautifulSoup 将其 解析为文本。我们可以通过传递参数来定制 HTML -> 文本解析 给 BeautifulSoup 解析器，使用 bs_kwargs (见 BeautifulSoup 文档). 在这种情况下，只有类为 “post-content”、“post-title” 或 “post-header” 的 HTML 标签是相关的，因此我们将移除所有其他标签。

"""
loader = WebBaseLoader(
    web_paths=("https://lilianweng.github.io/posts/2023-06-23-agent/",),
    bs_kwargs=dict(
        parse_only=bs4.SoupStrainer(
            class_=("post-content", "post-title", "post-header")
        )
    ),
)
docs = loader.load()

"""
我们加载的文档超过 42k 字符。这太长，无法适应 许多模型的上下文窗口。即使对于那些可以 将完整帖子放入其上下文窗口的模型，模型也可能难以找到 非常长输入中的信息。

为了解决这个问题，我们将把 Document 拆分成块以进行嵌入和 向量存储。这应该帮助我们在运行时仅检索最相关的部分 博客文章。

在这种情况下，我们将把文档拆分成每块1000个字符 并在块之间重叠200个字符。重叠有助于 减轻将陈述与相关重要 上下文分开的可能性。我们使用 RecursiveCharacterTextSplitter， 它将使用常见分隔符递归地拆分文档，例如 直到每个块的大小合适的新行。这是 通用文本用例的推荐文本分割器。

我们设置 add_start_index=True，以便每个 分割文档在初始文档中的字符索引得以保留，作为 元数据属性“start_index”。
"""
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
splits = text_splitter.split_documents(docs)

embeddings = DashScopeEmbeddings(
    model="text-embedding-v2",
    dashscope_api_key=os.getenv('LLM_API_KEY')
)
vectorstore = Chroma.from_documents(documents=splits, embedding=embeddings)

# Retrieve and generate using the relevant snippets of the blog.
retriever = vectorstore.as_retriever()
# 2. Incorporate the retriever into a question-answering chain.
system_prompt = (
    "You are an assistant for question-answering tasks. "
    "Use the following pieces of retrieved context to answer "
    "the question. If you don't know the answer, say that you "
    "don't know. Use three sentences maximum and keep the "
    "answer concise."
    "\n\n"
    "{context}"
)

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system_prompt),
        ("human", "{input}"),
    ]
)

question_answer_chain = create_stuff_documents_chain(llm, prompt)
rag_chain = create_retrieval_chain(retriever, question_answer_chain)
response = rag_chain.invoke({"input": "What is Task Decomposition?"})
print((response["answer"]))
